#
# Copyright 2008-2010 Vadim Zhukov <persgray@gmail.com>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY AGAVA COMPANY ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
# EVENT SHALL THE AGAVA COMPANY OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# The views and conclusions contained in the software and documentation are
# those of the authors and should not be interpreted as representing official
# policies, either expressed or implied, of the AGAVA Company.

package RB::Transfer;

use bytes;
use strict;

use Errno qw/EAGAIN EBADF EINVAL ETIMEDOUT/;
use Exporter qw/import/;
use MIME::Base64 qw/decode_base64 encode_base64/;

BEGIN {
	our @EXPORT = qw/bufferSize
	                 readWithTimeout readLineWithTimeout writeWithTimeout
	                 decodeHeaders encodeHeader endHeaders
	                 transferDataFromNet transferDataToNet
	                 sendArchive readArchives/;
}


# Buffer size, used by both dumper and agent
{ my $BufferSize = 65536;
sub bufferSize() {
	return $BufferSize;
} }

# Arguments:
#   $from    - file handle to read from
#   $timeout - timeout value (0 or undef means no timeout)
#   $size    - bytes ammount to read
#
# Returns:
#   - scalar buffer (where data read is put)
#
# Side effects:
#   Sets $! on error.
sub readWithTimeout(*$;$) {
	my ($from, $size, $timeout) = (shift, shift, shift);
	if ($timeout) {
		my $rin = '';
		vec($rin, fileno($from), 1) = 1;
		my $ein = $rin;
		my $ready = select($rin, undef, $ein, $timeout);
		if ($ready < 0) {
			return undef;
		} elsif ($ready == 0) {
			$! = ETIMEDOUT;
			return undef;
		}
	}

	my ($buf, $nread);
	unless (defined($nread = sysread($from, $buf, $size))) {
		return undef;
	} elsif ($nread == 0) {
		return '';
	}
	return $buf;
}

# Arguments:
#   $from    - file handle to read from
#   $timeout - timeout value (0 or undef means no timeout)
#
# Returns:
#   - text line (with or without trailing "\n")
#
# Side effects:
#   Sets $! on error.
sub readLineWithTimeout(*;$) {
	my ($from, $timeout) = (shift, shift);

	my $s = '';
	if ($timeout) {
		my ($rin, $bs) = ('', bufferSize);
		my $c;
		vec($rin, fileno($from), 1) = 1;
		while (1) {
			my ($rin2, $ein2) = ($rin, $rin);
			my $ready = select($rin2, undef, $ein2, $timeout);
			if ($ready < 0) {
				print STDERR "readLineWithTimeout(): ".
				    "select(): $!\n"
				    if ${main::Config}->{'test_mode'};
				return undef;
			} elsif ($ready == 0) {
				$! = ETIMEDOUT;
				print STDERR "readLineWithTimeout(): ".
				    "select(): $!\n";
				return undef;
			}
			my $rc = sysread($from, $c, 1);
			unless (defined $rc) {
				print STDERR "readLineWithTimeout(): ".
				    "sysread(): $!\n"
				    if ${main::Config}->{'test_mode'};
				return undef;
			} elsif ($rc == 0) {
				last;
			}
			$s .= $c;
			last if $c eq "\n";
		}
	} else {
		# Yes, it looks like ugly that we have to reimplement
		# readline() for unbuffered I/O. But tests show that speed
		# lowers for about 20%, which is, IMHO, not a big deal.
		# Especially when taking into account that this function
		# is used for reading (small) headers, not real (huge) data.
		while (1) {
			my $c;
			my $rc = sysread($from, $c, 1);
			unless (defined $rc) {
				print STDERR "readLineWithTimeout(): ".
				    "sysread(): $!\n"
				    if ${main::Config}->{'test_mode'};
				return undef;
			} elsif ($rc == 0) {
				last;
			}
			$s .= $c;
			last if $c eq "\n";
		}
	}
	return $s;
}

# Arguments:
#   $to      - file handle to write into
#   $timeout - timeout value (0 or undef means no timeout)
#   $buf     - scalar buffer to write
#
# Returns:
#   - ammount of data written, or undefined on error
#
# Side effects:
#   Sets $! on error.
sub writeWithTimeout(*$;$) {
	my ($to, $buf, $timeout) =
	    (shift, (shift || 0), shift);
	if ($timeout) {
		my $win = '';
		vec($win, fileno($to), 1) = 1;
		my $ein = $win;
		my $ready = select(undef, $win, $ein, $timeout);
		if ($ready < 0) {
			return undef;
		} elsif ($ready == 0) {
			$! = ETIMEDOUT;
			return undef;
		}
	}
	return syswrite($to, $buf);
}

# Arguments:
#   $line - encoded header string
#
# Returns:
#   - header name (scalar)
#   - header value (scalar)
#   - header attributes (hashref)
#   NOTE: returns empty list on error
#
# Side effects:
#   Sets $! on error.
sub decodeHeader($) {
	if($_[0] !~ m@^([A-Za-z_-]+)[\s]*:[\s]*(([A-Za-z0-9+/_]+)(([\s]+[^=\s]+=[A-Za-z0-9+/_]*)*))?$@) {
		print STDERR "Header string decoding failed: ".$_[0]."\n"
		    if ${main::Config}->{'test_mode'};
		$! = EINVAL;
		return ();
	}

	my ($name, $value, $attrStr) = (lc($1), $3, $4);
	my %attrs;
	$name =~ s/-/_/g;
	if ($value ne '') {
		$value =~ s/_/=/g;
		$value = decode_base64($value);

		while ($attrStr !~ /^[\s]*$/) {
			# We already checked all possible cases in regex above,
			# so do not bother about them much now.
			unless ($attrStr =~ /^[\s]+([^=\s]+)=([\S]*)/) {
				print STDERR "Header string attributes ".
				    "decoding failed: ".$_[0]."\n"
				    if ${main::Config}->{'test_mode'};
				$! = EINVAL;
				return ();
			} elsif (${main::Config}->{'test_mode'}) {
				print STDERR "Header attribute: $1 => $2\n";
			}
			$attrStr = substr($attrStr, $+[0]);
			my ($aname, $avalue) = (lc($1), $2);
			$avalue =~ s/_/=/g;
			$attrs{$aname} = decode_base64($avalue);
		}
	}
	return ($name, $value, \%attrs);
}

# Arguments:
#   $source - file handle to read headers from. Headers are read until empty
#             line discovered.
#
# Returns:
#   - structure: header->[]->{value, attributes};
#     attributes: attrname->attrvalue
#   NOTE: returns undefined value on error
#   NOTE: returns 0 on EOF
#
# Side effects:
#   Sets $! on error.
sub decodeHeaders(*) {
	my $from = shift;
	my %headers;
	my $timeout = ${main::Config}->{'network_timeout'};
	while (1) {
		my $headerStr = readLineWithTimeout($from, $timeout);
		unless (defined $headerStr) {
			print STDERR "decodeHeaders(): readLineWithTimeout() ".
			    "returned undef: $!\n"
			    if ${main::Config}->{'test_mode'};
			return undef;
		}
		if ($headerStr eq '') {
			print STDERR "decodeHeaders(): EOF detected when ".
			    (scalar(%headers) / 2)." headers got read\n"
			    if ${main::Config}->{'test_mode'};
			return %headers ? \%headers : 0;
		}
		chomp $headerStr;
		last if $headerStr eq '';    # end-of-headers marker
		my ($name, $value, $attrs) = decodeHeader $headerStr;
		unless (defined $name) {
			print STDERR "decodeHeaders(): illegal header string: ".
			    "$headerStr\n"
			    if ${main::Config}->{'test_mode'};
			$! = EINVAL;
			return undef;
		}
		$headers{$name} = [] unless exists $headers{$name};
		push(@{$headers{$name}}, { 'value' => $value,
		    'attributes' => $attrs });
	}
	return \%headers;
}

# Arguments:
#   $name       - header name (should not be empty)
#   $value      - header value (may be empty)
#   %attributes - pairs "name"->"value". May be omitted
#
# Returns:
#   - encoded header string
#
# Side effects:
#   Die on illegal parameter (header name is empty).
sub encodeHeader($;$%) {
	my ($name, $value) = (shift, shift);
	die "Empty header name" unless $name ne '';
	
	$name =~ s/_/-/g;
	$value = encode_base64($value, '');
	$value =~ s/=/_/g;
	my $s = $name.': '.$value;
	while (scalar(@_) >= 2) {
		my ($aname, $avalue) = (shift, shift);
		$avalue = encode_base64($avalue, '');
		$avalue =~ s/=/_/g;
		$s .= " ".$aname."=".$avalue;
	}
	return undef if @_;    # Illegal attribute declaration
	return $s."\n";
}

# Returns:
#   end-of-headers marker
sub endHeaders() {
	print STDERR "END-OF-HEADERS REQUESTED\n"
	    if ${main::Config}->{'test_mode'};
	return "\n";
}


# TODO: Treat read buffer size less than buffer size as error.
# TODO: Treat no separator/padding info at the end of data as error.

# Reads data block (archive) from network
#
# Arguments:
#   $from      - file handle to read data from
#   $to        - file handle to write into
#   $separator - separator to differentiate items
#
# Returns:
#   - boolean (true if no error, false if error occured)
#   - raw transferred bytes count
#   - data transferred bytes count
#
# Side effects:
#   - Sets $! on error.
sub transferDataFromNet(**$) {
	my ($from, $to, $separator) = splice(@_, 0, 3);
	my ($buf, $buf2);
	my ($bs, $buf1, $ok) = (bufferSize, '', 1);

	# Cache
	my $testmode = ${main::Config}->{'test_mode'};

# 	logIt 1, "Sending data";
	my $timeout = ${main::Config}->{'timeout'};

	# Use "lazy" reading, to have abillity to trim padding in last _data_
	# packet we receive. Padding size is 4-byte integer in network format
	# immediately following separator in last packet we receive.
	my ($bytesRaw, $bytesData) = (0, 0);
	while (1) {
		$buf = readWithTimeout($from, $bs - length($buf1), $timeout);
		unless (defined $buf) {
			$ok = 0;
			print STDERR "transferDataFromNet(): ".
			    "readWithTimeout(data) failed: $!\n"
			    if $testmode;
			last;
		}
		if (length($buf)) {
			$buf1 .= $buf;
			redo if length($buf1) < $bs;
			print STDERR "/" if $testmode;
			$bytesRaw += $bs;
		} else {
			$ok = 0;
			print STDERR "transferDataFromNet(): unexpected EOF\n"
			    if $testmode;
			$! = EAGAIN;
			last;
		}

		# Check for end-of-data marker
		if ($separator eq substr($buf1, 0, length($separator))) {
			print STDERR "SEPARATOR CAUGHT AT RAW $bytesRaw DATA $bytesData\n"
			    if $testmode;
			# Trim padding from last data buffer read
			my $padLen = unpack('N',
			    substr($buf1, length($separator), 4));
			$buf2 = substr($buf2, 0, $bs - $padLen);
			last;
		}

		# Do not try to write to undefined file handle
		next unless defined $to;

		if (defined $buf2) {
			# Abort on write error too
			unless (syswrite($to, $buf2)) {
				$ok = 0;
				print STDERR "transferDataFromNet(): ".
				    "print(data) failed: $!\n"
				    if $testmode;
				last;
			}
			$bytesData += length($buf2);
			print STDERR "\\" if $testmode;
		}
		$buf2 = $buf1;
	} continue {
		$buf1 = '';
	}
	# Do not check if error occured, it's connected to $buf1 anyway
	if (defined $buf2 and defined $to) {
		my $ec = $!+0;
		if (defined syswrite($to, $buf2)) {
			$bytesData += length($buf2);
			print STDERR "\\" if $testmode;
		} else {
			$ok = 0;
			print STDERR "transferDataFromNet(): ".
			    "print(tail) failed: $!\n"
			    if $testmode;
			$! = $ec if $ec;    # Do not mask first error by second
		}
	}
	print STDERR "\n" if $testmode;
	return ($ok, $bytesRaw, $bytesData);
}

# Sends data block (archive) over network
#
# Arguments:
#   $from      - file handle to read data from
#   $to        - file handle to write into
#   $separator - separator to differentiate items
#
# Returns:
#   - boolean (true if no error, false if error occured)
#   - raw transferred bytes count
#   - data transferred bytes count
#
# Side effects:
#   - Sets $! on error.
sub transferDataToNet(**$) {
	my ($from, $to, $separator) = splice(@_, 0, 3);
	my $buf;
	my ($rin, $win, $bs, $ok) = ('', '', bufferSize, 1);

	# Cache
	my $testmode = ${main::Config}->{'test_mode'};
	my $timeout = ${main::Config}->{'timeout'};
	
	my ($bytesRaw, $bytesData) = (0, 0);
	while (1) {
		my $rc = sysread($from, $buf, $bs);
		unless (defined $rc) {
			$ok = 0;
			print STDERR "transferDataToNet(): ".
			    "sysread() failed: $!\n"
			    if $testmode;
			last;
		} elsif ($rc == 0) {
			print STDERR
			    "transferDataToNet(): EOF, $bytesRaw bytes read\n"
			    if $testmode;
			# Normal EOF
			last;
		}
		$bytesRaw += length($buf);

		# Do not try to write to undefined file handled
		if (defined $to) {
			# Abort on write error too
			unless (writeWithTimeout($to, $buf, $timeout)) {
				$ok = 0;
				print STDERR "transferDataToNet(): ".
				    "writeWithTimeout(data) failed: $!\n"
				    if $testmode;
				last;
			}
			print STDERR "+" if $testmode;
		} else {
			print STDERR "+" if $testmode;
		}
	}
	$bytesData = $bytesRaw;
	if ($ok) {
		my $rem = $bytesRaw % $bs;
		my $padLen = 0;
		if ($bytesData == 0) {
			# We have to fill one, even empty, block always.
			$padLen = $bs;
		} elsif ($rem > 0) {
			# Padding, makes searching for separator easy.
			$padLen = $bs - $rem unless $bytesData == 0;
		}
		if ($padLen) {
			print STDERR "PADDING IS $padLen BYTES\n"
			    if $testmode;
			unless (writeWithTimeout($to, pack('x'.$padLen),
			    $timeout)) {
				$ok = 0;
				print STDERR "transferDataToNet(): ".
				    "writeWithTimeout(padding) failed: $!\n"
				    if $testmode;
				goto TRANSFER_TO_NET_END;
			}
		}
		$bytesRaw += $padLen;

		# Send just the size of padding, to avoid integer overflow
		# NOTE: pad this too. Yes, we lose almost 64KB for nothing :(
		my $footerPadLen = $bs - (length($separator) + 4);
		if (writeWithTimeout($to,
		    $separator.pack('Nx'.$footerPadLen, $padLen), $timeout)) {
			$bytesRaw += $bs;
			print STDERR "+" if $testmode;
		} else {
			$ok = 0;
			print STDERR "transferDataToNet(): ".
			    "writeWithTimeout(footer) failed: $!\n"
			    if $testmode;
			goto TRANSFER_TO_NET_END;
		}
	}
TRANSFER_TO_NET_END:
	print STDERR "\n" if $testmode;
	print STDERR "TRANSFERINg TO NET: ".($ok ? "OK" : "FAIL")."\n"
	    if $testmode;
	return ($ok, $bytesRaw, $bytesData);
}

# Sends "header" and archive itself.
#
# Arguments:
#   $from      - file handle to read data from
#   $to        - file handle to write into
#   $name      - archive (item) name
#   $separator - separator to differentiate items
#
# Returns:
#   - boolean (true if no error, false if error occured)
#   - raw transferred bytes count
#   - data transferred bytes count
#
# Side effects:
#   - Sets $! on error.
#   - Assumes that it is run in "headers" mode, and do not send any headers
#     after transfer itself.
sub sendArchive(**$$) {
	my ($from, $to, $name, $separator) = splice(@_, 0, 4);

	my ($bs, $timeout) = (bufferSize, ${main::Config}->{'network_timeout'});
	my $header = encodeHeader('sending_item', $name).endHeaders();

	unless (writeWithTimeout($to, $header, $timeout)) {
		print STDERR "sendArchive(): writeWithTimeout(header)".
		    " returned error: $!\n"
		    if ${main::Config}->{'test_mode'};
		return (0, 0, 0);
	}
	print STDERR "SENDARCHIVE: HEADERS SENT\n"
	    if ${main::Config}->{'test_mode'};
	my @ret = transferDataToNet($from, $to, $separator);
	return wantarray ? @ret : $ret[0];
}

# Arguments:
#   $from      - file handle to read data from
#   $separator - separator to differentiate items
#   $hookOpen  - reference to function called to open destination filehandle
#   $hookClose - reference to function called to close destination filehandle.
#   Subsequent parameters are passed to hooks unchanged
#
# Returns:
#   - arrayref of item structures, in order of dump.
#     NOTE: At least one item (with undefined name in case of early error)
#           already presents there.
#     NOTE: "headers" field is always valid hashref.
#   - raw transferred bytes count
#   - data transferred bytes count
#
# It's guaranteed that every $hookOpen call will be paired with $hookClose call.
sub readArchives($$&&@) {
	my ($from, $separator, $hookOpen, $hookClose) = splice(@_, 0, 4);
	my ($wholeBytesRaw, $wholeBytesData, $closed) = (0, 0, 1);
	my (@downloadedItems, $item, $to, $h);
	my $testmode = ${main::Config}->{'test_mode'};
	my $listDumped = ${main::Config}->{'list_dumped'};
	while (1) {
		print STDERR "DECODING PRE-HEADERS\n"
		    if $testmode;
		$h = decodeHeaders($from);
		$item = { 'name' => undef, 'headers' => $h };
		unless (defined $h) {
			$item->{'errcode'} = $!+0 || EBADF;
			$item->{'error'} = 'premature end of headers';
			$item->{'headers'} = {};
			push(@downloadedItems, $item);
			last;
		}
		last if $h == 0;    # EOF
		push(@downloadedItems, $item);

		if ($testmode) {
			foreach my $_hn (sort keys %{$h}) {
				my $_h = $h->{$_hn};
				foreach my $i (sort { $_h->[$a]->{'value'}
				    cmp $_h->[$b]->{'value'} } 0..$#{$_h}) {
					print STDERR "READ HEADER: ".
					    $_hn.": ".$_h->[$i]->{'value'};
					foreach my $_an (sort keys
					    %{$_h->[$i]->{'attributes'}}) {
						print STDERR " ".$_an."=".
						    $_h->[$i]->{'attributes'}->
						    {$_an};
					}
					print STDERR "\n";
				}
			}
		}

		unless ($h->{'sending_item'}) {
			$item->{'errcode'} = EBADF;
			$item->{'error'} = 'no item name sent';
			last;
		}
		$item->{'name'} = $h->{'sending_item'}->[0]->{'value'};
		print "dumping ".$item->{'name'}."\n"
		    if $listDumped;

		$to = &{$hookOpen}($item, @_);
		$closed = !defined($to);
		my ($ok, $bytesRaw, $bytesData) =
		    transferDataFromNet($from, $to, $separator);
		$wholeBytesRaw += $bytesRaw;
		$wholeBytesData += $bytesData;
		$item->{'transferred'} = $bytesData;
		unless ($ok) {
			$item->{'errcode'} = $!+0;
			$item->{'error'} = "transferring data: $!";
			last;
		}

		print STDERR "DECODING POST-HEADERS\n"
		    if $testmode;
		$h = decodeHeaders($from);
		unless (defined $h and $h != 0) {
			$item->{'errcode'} = EBADF;
			$item->{'error'} = 'no finalizing headers';
			last;
		}
		$item->{'errcode'} = 0;
		$item->{'headers'} = { %{$item->{'headers'}}, %{$h} };
		$closed = 1;    # We assume hook will close it anyway
		last unless &{$hookClose}($item, $to, @_);
	}
	unless ($closed) {
		&{$hookClose}($item, $to, @_);
	}
	unless (@downloadedItems) {
		@downloadedItems = ({
		    'name' => undef,
		    'errcode' => EAGAIN,
		    'error' => "connecting to agent: $!"
		    });
	}

	return (\@downloadedItems, $wholeBytesRaw, $wholeBytesData);
}


1;
